import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone
from collections import Counter

# 初始化 Pinecone
# 请确保你的 API 密钥和索引名称是正确的
pinecone = Pinecone(api_key="pcsk_31YZA5_GanCvpQrrdy6qYfssjLYpt7J3n5R7t84Jt9CWmLzog21cCyireQLRj3hd6oWq58") # 替换成你的 Pinecone API Key
index_name = "mnist-index"

# 检查索引是否存在，如果不存在则提示
if index_name not in pinecone.list_indexes().names():
    raise ValueError(f"Pinecone index '{index_name}' does not exist. Please create it first.")

# 连接到 Pinecone 索引
index = pinecone.Index(index_name)

# 1. 预处理图像函数 (保持不变)
#    这个函数负责将用户绘制的图像转换为模型可以理解的格式 (一个 64 维的向量)
def preprocess(image):
    # 将输入的NumPy数组转换为PIL图像对象
    image = Image.fromarray(image)
    # 调整大小为 8x8，并转换为灰度图
    image = image.resize((8, 8)).convert('L')
    # 将图像转换为 NumPy 数组
    image_array = np.array(image)
    # 将 8x8 的二维数组展平为 64 个元素的一维数组
    img_vector = image_array.ravel()
    return img_vector

# 2. 模型预测函数 (稍作简化)
#    这个函数只负责查询 Pinecone，输入参数是已经预处理好的向量
def model_predict(img_vector):
    # 将 NumPy 数组转换为 Python 列表，以符合 Pinecone 的查询要求
    img_list = img_vector.tolist()

    # 使用准备好的查询向量在 Pinecone 索引中执行搜索
    # top_k=11 表示我们寻找 11 个最相似的邻居
    results = index.query(
        vector=img_list,
        top_k=11,
        include_metadata=True
    )

    # 从搜索结果中提取每个匹配项的标签 (即它们代表的数字)
    labels = [match['metadata']['label'] for match in results['matches']]

    # 使用投票机制（少数服从多数）来确定最终的分类结果
    # Counter(labels).most_common(1) 会找到出现次数最多的标签
    final_prediction = Counter(labels).most_common(1)[0][0]
    return final_prediction

# 3. 创建一个新的“流水线”函数 (这是关键的修改)
#    这个函数连接了预处理和预测两个步骤
def recognize_digit(sketchpad_image):
    # 第一步：调用 preprocess 函数处理原始图像
    processed_vector = preprocess(sketchpad_image)
    # 第二步：将处理好的向量传递给 model_predict 函数进行预测
    prediction = model_predict(processed_vector)
    # 返回最终的预测结果
    return prediction

# 创建 Gradio 界面
iface = gr.Interface(
    fn=recognize_digit,  # <-- 关键改动：让界面调用我们新的流水线函数
    inputs=gr.Sketchpad(label="请在此处绘制一个 0-9 的数字", type="numpy"),
    outputs=gr.Label(num_top_classes=1, label="预测结果"),
    title="k-NN 手写数字识别",
    description="在左侧画板上用鼠标绘制一个数字，模型将使用 k-NN 算法（通过 Pinecone）进行识别。"
)

# 启动 Gradio 应用
iface.launch(share=True)